# -*- coding: utf-8 -*-
'''
Created on Jul 21, 2019

@author: yl
'''
from FFT_Interpolation import *
from mpl_toolkits.mplot3d import Axes3D
from scipy import signal
from scipy.optimize import curve_fit
from scipy.signal import *
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation

fig = plt.figure()


j = complex(0, 1)
c = 3e8 # 光速 [m/s]
Lamda = 633e-9 # 光波长 [m]
Fc = c / Lamda # 光频率 [Hz]

pix_size = 5.3e-6
pix_num = 1280
screen_diameter = pix_num * pix_size


def f(X, Y):
    window = signal.gaussian(pix_num, std=pix_num/10)
    window_x, window_y = np.meshgrid(window, window)
    window_2d = window_x*window_y
     
    I_0 = 220
    D_d = D - X * np.tan(V_x_2 - V_x) + Y * np.tan(V_y_2 - V_y)
    L_d = L + X * V_x + Y * V_y
    diff_L_d_2 = D_d + (L_d + D_d) * (1 - (V_x_2**2 + V_y_2**2)) / (1 + (V_x_2**2 + V_y_2**2)) - L_d * (1 - (V_x**2 + V_y**2)) / (1 + (V_x**2 + V_y**2))
    Z = I_0 * (1 + np.cos(2 * np.pi * diff_L_d_2 / Lamda))
     
    R = 1
    center = [(pix_num-1+800)*pix_size/2, (pix_num-800)*pix_size/2]
    r = np.sqrt((X-center[0])**2 + (Y-center[1])**2)
    d = R - np.sqrt(R**2 - r**2)
    A = 2 * np.pi * (d+2*D) / Lamda
    I_r = np.sin(A)**2 * I_0/8
    
    R = 2
    center = [(pix_num-1-800)*pix_size/2, (pix_num+800)*pix_size/2]
    r = np.sqrt((X-center[0])**2 + (Y-center[1])**2)
    d = R - np.sqrt(R**2 - r**2)
    A = 2 * np.pi * (d+2*D) / Lamda
    I_r_2 = np.sin(A)**2 * I_0/2
    
    return (Z+I_r_2)*window_2d


# def f(X, Y):
#     R = 0.1
#     center = [(pix_num-1)*pix_size/2, (pix_num-1)*pix_size/2]
#     r = np.sqrt((X-center[0])**2 + (Y-center[1])**2)
#     d = R - np.sqrt(R**2 - r**2)
#     A = 2 * np.pi * (d+2*D) / Lamda
#     I_r = np.sin(A)**2
#     return I_r

dx = np.linspace(0, (pix_num-1)*pix_size, num=pix_num)
dy = dx
V_x, V_y, Vz = 0.000, 0.000, 1 # Reference
V_x_2, V_y_2, Vz_2 = 0.0005, 0.000, 1 # Measurement
tana = (V_x**2 + V_y**2)**0.5
tana_2 = (V_x_2**2 + V_y_2**2)**0.5
L, D = 20, 200
X, Y = np.meshgrid(dx, dy)

# plt.subplot(2,1,1)
im = plt.imshow(f(X, Y),cmap='gray', animated=True)
# plt.colorbar(im, shrink=0.6)
plt.colorbar(im, fraction=0.046, pad=0.04)

# plt.subplot(2,1,2)
# line = plt.plot(f(X, Y)[640], animated=True)
line, = plt.plot((-f(X, Y)[640]/3+1000), color='white')

def updatefig(*args):
    global D
    D += -Lamda / 16.
    im.set_array(f(X, Y))
    line.set_ydata(-f(X, Y)[640]/3+1200)
    return im, line,

ani = animation.FuncAnimation(fig, updatefig, interval=50, blit=True)
plt.show()
# ani.save('test.gif', writer=animation.PillowWriter(fps=24))
# ani.save('test.gif')